"""
震动传播模型

实现了横向和纵向的震动传播模型，包括：
- 横向传播（从左向右）
- 纵向传播（从下向上）
- 传播速度控制
- 步长控制
"""

import numpy as np
from typing import Literal, Tuple, Optional
from loguru import logger
from .decay_model import DecayModel


class VibrationPropagation:
    """震动传播模型"""
    
    def __init__(self, 
                 direction: Literal["horizontal", "vertical"] = "horizontal",
                 propagation_speed: float = 1000.0,
                 step_size: float = 0.01,
                 attenuation_coefficient: float = 0.1):
        """
        初始化震动传播模型
        
        Args:
            direction: 传播方向，horizontal为横向，vertical为纵向
            propagation_speed: 传播速度 (m/s)
            step_size: 传播步长 (m)
            attenuation_coefficient: 衰减系数
        """
        self.direction = direction
        self.propagation_speed = propagation_speed
        self.step_size = step_size
        self.attenuation_coefficient = attenuation_coefficient
        
        # 初始化衰减模型
        self.decay_model = DecayModel(decay_type="exponential")
        
        logger.info(f"初始化{direction}震动传播模型，速度={propagation_speed}m/s，步长={step_size}m")
    
    def propagate_vibration(self, 
                          initial_amplitude: float,
                          source_position: Tuple[float, float],
                          target_positions: np.ndarray,
                          time_steps: int = 100) -> np.ndarray:
        """
        计算震动传播到目标位置的振幅
        
        Args:
            initial_amplitude: 初始振幅
            source_position: 震源位置 (x, y)
            target_positions: 目标位置数组，形状为 (n, 2)
            time_steps: 时间步数
            
        Returns:
            目标位置的振幅数组，形状为 (n, time_steps)
        """
        n_targets = len(target_positions)
        amplitudes = np.zeros((n_targets, time_steps))
        
        for i, target_pos in enumerate(target_positions):
            distance = self._calculate_distance(source_position, target_pos)
            amplitudes[i, :] = self.decay_model.calculate_amplitude(
                initial_amplitude, distance, self.attenuation_coefficient
            )
        
        return amplitudes
    
    def simulate_propagation(self, 
                           initial_amplitude: float,
                           source_position: Tuple[float, float],
                           grid_size: Tuple[int, int] = (50, 50),
                           grid_spacing: float = 0.01,
                           time_steps: int = 100) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        模拟震动在网格上的传播
        
        Args:
            initial_amplitude: 初始振幅
            source_position: 震源位置 (x, y)
            grid_size: 网格大小 (rows, cols)
            grid_spacing: 网格间距 (m)
            time_steps: 时间步数
            
        Returns:
            (x坐标网格, y坐标网格, 振幅网格)
        """
        rows, cols = grid_size
        
        # 创建坐标网格
        x = np.arange(0, cols * grid_spacing, grid_spacing)
        y = np.arange(0, rows * grid_spacing, grid_spacing)
        X, Y = np.meshgrid(x, y)
        
        # 计算每个网格点到震源的距离
        distances = np.sqrt((X - source_position[0])**2 + (Y - source_position[1])**2)
        
        # 计算振幅
        amplitudes = self.decay_model.calculate_amplitude(
            initial_amplitude, distances, self.attenuation_coefficient
        )
        
        return X, Y, amplitudes
    
    def _calculate_distance(self, 
                          source_pos: Tuple[float, float], 
                          target_pos: np.ndarray) -> float:
        """
        计算震源到目标位置的距离
        
        Args:
            source_pos: 震源位置 (x, y)
            target_pos: 目标位置 (x, y)
            
        Returns:
            距离
        """
        return np.sqrt((target_pos[0] - source_pos[0])**2 + (target_pos[1] - source_pos[1])**2)
    
    def get_propagation_time(self, distance: float) -> float:
        """
        计算震动传播时间
        
        Args:
            distance: 传播距离
            
        Returns:
            传播时间 (s)
        """
        return distance / self.propagation_speed
    
    def plot_propagation(self, 
                        initial_amplitude: float,
                        source_position: Tuple[float, float],
                        grid_size: Tuple[int, int] = (50, 50),
                        grid_spacing: float = 0.01,
                        title: str = None):
        """
        绘制震动传播图
        
        Args:
            initial_amplitude: 初始振幅
            source_position: 震源位置
            grid_size: 网格大小
            grid_spacing: 网格间距
            title: 图表标题
        """
        try:
            import matplotlib.pyplot as plt
            
            X, Y, amplitudes = self.simulate_propagation(
                initial_amplitude, source_position, grid_size, grid_spacing
            )
            
            plt.figure(figsize=(12, 8))
            
            # 绘制振幅等高线图
            contour = plt.contourf(X, Y, amplitudes, levels=20, cmap='viridis')
            plt.colorbar(contour, label='振幅')
            
            # 标记震源位置
            plt.plot(source_position[0], source_position[1], 'r*', 
                    markersize=15, label='震源')
            
            plt.xlabel('X 坐标 (m)')
            plt.ylabel('Y 坐标 (m)')
            plt.title(title or f'{self.direction}震动传播')
            plt.legend()
            plt.grid(True, alpha=0.3)
            plt.axis('equal')
            plt.show()
            
        except ImportError:
            logger.warning("matplotlib未安装，无法绘制传播图")
    
    def get_propagation_parameters(self) -> dict:
        """
        获取传播参数
        
        Returns:
            参数字典
        """
        return {
            'direction': self.direction,
            'propagation_speed': self.propagation_speed,
            'step_size': self.step_size,
            'attenuation_coefficient': self.attenuation_coefficient
        }
    
    def set_propagation_parameters(self, **kwargs):
        """
        设置传播参数
        
        Args:
            **kwargs: 要设置的参数
        """
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
                logger.info(f"设置{key}={value}")
            else:
                logger.warning(f"未知参数: {key}")
    
    def __str__(self) -> str:
        return (f"VibrationPropagation(direction={self.direction}, "
                f"speed={self.propagation_speed}m/s, step={self.step_size}m)")
    
    def __repr__(self) -> str:
        return self.__str__() 